from collections import deque
from common.tools import preload_time
from common.variables import time_counts, TimeFrameType
from trade.data.add_factors import add_batch_factors
from trade.data.table_edit import get_factor_from_table_by_pivot
from trade.data.watch_kline_data import WatchFactorCenter
from trade.factors import name_bfname


RecordsType = dict[tuple[str, str], deque[float]]


class FactorLooker:
    data_records: RecordsType
    factor_gather: WatchFactorCenter

    def __init__(
        self,
        need_symbols: list[str],
        need_fnames: list[str],
        timeframe: TimeFrameType,
    ):
        self.data_records = {}
        self.need_symbols = need_symbols
        self.need_fnames = need_fnames
        self.timeframe = timeframe
        self.factor_gather = WatchFactorCenter(need_symbols, need_fnames, timeframe)

    @staticmethod
    def get_symbol_and_factors(
        start: int,
        end: int,
        need_symbols: list[str],
        need_fnames: list[str],
        timeframe: TimeFrameType,
    ):
        data_records = {}
        if not need_symbols:
            return data_records

        data_records: RecordsType
        all_bfnames = list(set([name_bfname[x] for x in need_fnames]))
        add_batch_factors(
            all_bfnames,
            need_symbols,
            start,
            end,
            timeframe,
            max_part=2000,
            n_jobs=1,
        )

        for symbol in need_symbols:

            df = get_factor_from_table_by_pivot(
                timeframe,
                start,
                end,
                symbol,
                need_fnames,
            )
            if (key := ("default", "xdatetime")) not in data_records:
                data_records[key] = df.index

            for col in df.columns:
                data_records[(symbol, col)] = deque(df[col], maxlen=len(df[col]))

        sizes = [len(x) for x in data_records.values()]
        if not sizes:
            return data_records

        if not all([x == sizes[0] for x in sizes]):
            raise Exception(
                f"初始化数据缺失：{start} {end} {need_symbols} {need_fnames}"
            )

        return data_records

    def db_for_factors(
        self,
        start: int,
        end: int,
        preloads: int,
        nsize: int,
    ):
        timeframe = self.timeframe
        need_symbols = self.need_symbols
        need_fnames = self.need_fnames

        data_records = self.get_symbol_and_factors(
            preload_time(start, timeframe, preloads),
            start,
            self.need_symbols,
            self.need_fnames,
            self.timeframe,
        )
        self.data_records = data_records

        data_records = self.get_symbol_and_factors(
            start,
            preload_time(start, timeframe, -nsize),
            need_symbols,
            need_fnames,
            timeframe,
        )

        while data_records[("default", "xdatetime")]:
            if self.data_records[("default", "xdatetime")][-1] > end:
                raise StopIteration
            for kk in data_records:
                self.data_records[kk].popleft()
                self.data_records[kk].append(data_records[kk].popleft())
            yield
            if not data_records[("default", "xdatetime")]:
                p = (
                    self.data_records[("default", "xdatetime")][-1]
                    + time_counts[timeframe] * 1000
                )
                data_records = FactorLooker.get_symbol_and_factors(
                    p,
                    preload_time(p, timeframe, preloads),
                    need_symbols,
                    need_fnames,
                    timeframe,
                )
        raise StopIteration(f"循环遍历数据，没有数据了")

    def online_for_factors(self, preloads: int):
        tmp = self.factor_gather.get_batch_factors()
        start = tmp[("default", "xdatetime")]
        data_records = self.get_symbol_and_factors(
            preload_time(start, self.timeframe, preloads),
            start,
            self.need_symbols,
            self.need_fnames,
            self.timeframe,
        )
        self.data_records = data_records
        while rr := self.factor_gather.get_batch_factors():
            for kk in rr:
                self.data_records[kk].popleft()
                self.data_records[kk].append(rr[kk])
            yield
